#include "broadcast_helper.h"
#include "kernel_operator.h"
#include "op_common.h"
#include <cstdint>
#include <type_traits>

using namespace AscendC;

template <typename _DT_START, typename _DT_END, typename _DT_WEIGHT,
          typename _DT_Y>
class Kernel {
public:
  TPipe pipe;
  CommonTiling tilingData;
  BroadcastHelper broadcastHelpr;
  ShapeData startShape;
  ShapeData endShape;
  ShapeData weightShape;
  DefInTensor(START);
  DefInTensor(END);
  DefInTensor(WEIGHT);
  DefOutTensor(Y);

  DefBufVECIN(TEMP1);
  DefBufVECIN(TEMP2);
  DefBufVECIN(HALFMAX);
  DefBufVECIN(HALFMIN);

public:
  __aicore__ inline Kernel() {}

  __aicore__ inline void Init(GM_ADDR start, GM_ADDR end, GM_ADDR weight,
                              GM_ADDR y, CommonTiling &tiling,
                              int64_t *broadcast_shape, int64_t *start_shape,
                              int64_t *end_shape, int64_t *weight_shape) {
    tilingData = tiling;
    broadcastHelpr.broadcastShape.init(broadcast_shape);
    startShape.init(start_shape);
    endShape.init(end_shape);
    weightShape.init(weight_shape);

    // 切分global
    GSetBuffer(START, start, 0, startShape.size);
    GSetBuffer(END, end, 0, endShape.size);
    GSetBuffer(WEIGHT, weight, 0, weightShape.size);
    GSetBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);

    if (!startShape.needBroadBast && !endShape.needBroadBast) {
      // 初始化队列
      InitQueueSimple(START, tilingData.bufferSize);
      InitQueueSimple(END, tilingData.bufferSize);
      InitQueueSimple(WEIGHT, tilingData.bufferSize);
      InitQueueSimple(Y, tilingData.bufferSize);
      // 初始化BUF
      if constexpr (std::is_same_v<TypeOf(START), half>) {
        InitTBufBuffer(TEMP1, tilingData.bufferSize * sizeof(float));
        InitTBufBuffer(TEMP2, tilingData.bufferSize * sizeof(float));
        InitTBufBuffer(HALFMAX, tilingData.bufferSize * sizeof(float));
        InitTBufBuffer(HALFMIN, tilingData.bufferSize * sizeof(float));
        TBufGet(HALFMAX, int32_t);
        TBufGet(HALFMIN, int32_t);
        //   0x477fe000 65504
        Muls(BTensor(HALFMAX), BTensor(HALFMAX), int32_t(0),
             tilingData.bufferSize);
        Adds(BTensor(HALFMAX), BTensor(HALFMAX), int32_t(0x477fe000),
             tilingData.bufferSize);
        //  -0x38802000 -65504
        Muls(BTensor(HALFMIN), BTensor(HALFMIN), int32_t(0),
             tilingData.bufferSize);
        Adds(BTensor(HALFMIN), BTensor(HALFMIN), int32_t(-0x38802000),
             tilingData.bufferSize);
      }
    }
  }

  __aicore__ inline void Process() {
    if (!startShape.needBroadBast && !endShape.needBroadBast) {
      auto finnal_progress = tilingData.loopCount - 1;
      for (uint32_t i = 0; i < finnal_progress; ++i) {
        CopyIn(i, tilingData.bufferSize);
        Compute(tilingData.bufferSize);
        CopyOut(i, tilingData.bufferSize);
      }
      CopyIn(finnal_progress, tilingData.finalCalcCount);
      Compute(tilingData.finalCalcCount);
      CopyOutFinal(finnal_progress, tilingData.finalCalcCount);
    } else {
      for (auto i = 0; i < broadcastHelpr.broadcastShape.size; ++i) {
        auto startOffset =
            broadcastHelpr.broadcastIndex2targetOffset(startShape);
        auto endOffset = broadcastHelpr.broadcastIndex2targetOffset(endShape);
        auto weightOffset =
            broadcastHelpr.broadcastIndex2targetOffset(weightShape);
        if constexpr (std::is_same_v<_DT_START, half>) {
          auto startValue = float(GTensor(START)(startOffset));
          auto v =
              startValue + float(GTensor(WEIGHT)(weightOffset)) *
                               (float(GTensor(END)(endOffset)) - startValue);
          if (v > 65504.0f) {
            GTensor(Y)(i) = infHalf;
          } else if (v < -65504.0f) {
            GTensor(Y)(i) = ninfHalf;
          } else {
            GTensor(Y)(i) = v;
          }
        } else {
          auto startValue = GTensor(START)(startOffset);
          GTensor(Y)(i) =
              startValue + GTensor(WEIGHT)(weightOffset) *
                               (GTensor(END)(endOffset) - startValue);
        }
        broadcastHelpr.broadcastIndexIncrease(1);
      }
    }
  }

  __aicore__ inline void CopyIn(uint32_t i, uint32_t calcCount) {
    int64_t startOffset = tilingData.startIdx + i * tilingData.bufferSize;
    auto copyCount = ALIGN_TO(calcCount, 32 / sizeof(TypeOf(START)));
    QueAllocSimple(START);
    if (startShape.needBroadBast) {
      broadcastHelpr.setIndexByOffset(startOffset);
      broadcastHelpr.DataCopy(LTensorName(START), GTensorName(START),
                              startShape, calcCount);
    } else {
      DataCopy(LTensorName(START), GTensorName(START)[startOffset], copyCount);
    }
    EnQue(START);
    QueAllocSimple(END);
    if (endShape.needBroadBast) {
      broadcastHelpr.setIndexByOffset(startOffset);
      broadcastHelpr.DataCopy(LTensorName(END), GTensorName(END), endShape,
                              calcCount);
    } else {
      DataCopy(LTensorName(END), GTensorName(END)[startOffset], copyCount);
    }
    EnQue(END);
    QueAllocSimple(WEIGHT);
    if (weightShape.needBroadBast) {
      broadcastHelpr.setIndexByOffset(startOffset);
      broadcastHelpr.DataCopy(LTensorName(WEIGHT), GTensorName(WEIGHT),
                              weightShape, calcCount);
    } else {
      DataCopy(LTensorName(WEIGHT), GTensorName(WEIGHT)[startOffset],
               copyCount);
    }
    EnQue(WEIGHT);
  }

  __aicore__ inline void Compute(uint32_t calcCount) {
    if constexpr (std::is_same_v<TypeOf(START), half>) {
      ComputeHalf(calcCount);
      return;
    }

    DeQueSimple(START);
    DeQueSimple(END);
    DeQueSimple(WEIGHT);
    QueAllocSimple(Y);
    // y = start + weight * (end - start)
    Sub(LTensorName(END), LTensorName(END), LTensorName(START), calcCount);
    Mul(LTensorName(END), LTensorName(END), LTensorName(WEIGHT), calcCount);
    Add(LTensorName(Y), LTensorName(START), LTensorName(END), calcCount);
    EnQue(Y);
    QueFree(START);
    QueFree(END);
    QueFree(WEIGHT);
  }

  __aicore__ inline void ComputeHalf(uint32_t calcCount) {
    auto cmpCount = ALIGN_TO(calcCount, elementsPerRepeat<float>());
    DeQueSimple(START);
    DeQueSimple(END);
    DeQueSimple(WEIGHT);
    QueAllocSimple(Y);
    TBufGet(TEMP1, float);
    TBufGet(TEMP2, float);
    TBufGet(HALFMIN, float);
    TBufGet(HALFMAX, float);
    // y = start + weight * (end - start)
    Cast(TBufTensorName(TEMP1), LTensorName(START), RoundMode::CAST_NONE,
         calcCount);
    Cast(TBufTensorName(TEMP2), LTensorName(END), RoundMode::CAST_NONE,
         calcCount);
    Sub(TBufTensorName(TEMP2), TBufTensorName(TEMP2), TBufTensorName(TEMP1),
        calcCount);

    Cast(TBufTensorName(TEMP1), LTensorName(WEIGHT), RoundMode::CAST_NONE,
         calcCount);
    Mul(TBufTensorName(TEMP2), TBufTensorName(TEMP2), TBufTensorName(TEMP1),
        calcCount);

    Cast(TBufTensorName(TEMP1), LTensorName(START), RoundMode::CAST_NONE,
         calcCount);
    Add(TBufTensorName(TEMP2), TBufTensorName(TEMP2), TBufTensorName(TEMP1),
        calcCount);

    Cast(LTensorName(Y), TBufTensorName(TEMP2), RoundMode::CAST_NONE,
         calcCount);

    auto cmpRet = LTensor(START).template ReinterpretCast<int8_t>();
    Compare(cmpRet, BTensor(TEMP2), BTensor(HALFMAX), CMPMODE::LE, cmpCount);
    Select(LTensor(Y), cmpRet, LTensor(Y), infHalf,
           SELMODE::VSEL_TENSOR_SCALAR_MODE, calcCount);
    Compare(cmpRet, BTensor(TEMP2), BTensor(HALFMIN), CMPMODE::GE, cmpCount);
    Select(LTensor(Y), cmpRet, LTensor(Y), ninfHalf,
           SELMODE::VSEL_TENSOR_SCALAR_MODE, calcCount);
    EnQue(Y);
    QueFree(START);
    QueFree(END);
    QueFree(WEIGHT);
  }

  __aicore__ inline void CopyOut(uint32_t i, uint32_t calcCount) {
    DeQueLocal2Global(Y, i * tilingData.bufferSize, calcCount);
  }

  __aicore__ inline void CopyOutFinal(uint32_t i, uint32_t calcCount) {
    int64_t startOffset = tilingData.startIdx + i * tilingData.bufferSize;
    DeQueSimple(Y);
    myDataCopy(GTensorName(Y), LTensorName(Y), calcCount,
               i * tilingData.bufferSize);
    QueFree(Y);
  }
};
extern "C" __global__ __aicore__ void lerp(GM_ADDR start, GM_ADDR end,
                                           GM_ADDR weight, GM_ADDR y,
                                           GM_ADDR workspace, GM_ADDR tiling) {
  GET_TILING_DATA(tiling_data, tiling);
  CRATE_COMMON_TILING_SIMPLE;
  Kernel<DTYPE_START, DTYPE_END, DTYPE_WEIGHT, DTYPE_Y> op;
  op.Init(start, end, weight, y, commonTiling, tiling_data.broadcastShape,
          tiling_data.startShape, tiling_data.endShape,
          tiling_data.weightShape);
  op.Process();
}